/*
Copyright 2025 PerfXLab (Beijing) Technologies Co., Ltd.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/
#include "OpenSparseBLAS_func.h"
#include <complex>
#include <vector>


namespace _spblasinl_spsm {
    template <typename T>
    inline T abs_val(const T& v) { return std::fabs(v); }
    template <typename T>
    inline T abs_val(const std::complex<T>& v) { return std::abs(v); }

    /* 安全下标 (仅 debug 用，可改 assert) */
    inline bool in_range(int idx, int len) { return idx >= 0 && idx < len; }

    /* 指针合法性简单检查 */
    template <typename... Ptrs>
    inline bool all_nonnull(Ptrs... ps) {
        return (... && (ps != nullptr));
    }
} 


/*****************  CSC sparse triangular solve: multiple RHS  *****************/
template <typename T>
OpenSpB_status OpenSpB_spsm(  int m, 
                                    int nrhs,
                                    const int *csc_col_ptr,
                                    const int *csc_row_ind,
                                    const T   *csc_vals,
                                    T *X, 
                                    int ldx,
                                    const T *B, 
                                    int ldb,
                                    OpenSpB_base_type index_base,
                                    OpenSpB_uplo_type uplo)
{
    using namespace _spblasinl_spsm;

    if (m<=0 || nrhs<=0) return OPENSPBLAS_STATUS_SUCCESS;
    if (!all_nonnull(csc_col_ptr,csc_row_ind,csc_vals,X,B))
        return OPENSPBLAS_STATUS_INVALID;
    if (ldx < m || ldb < m)   return OPENSPBLAS_STATUS_INVALID;

    const int base = (index_base==OPENSPBLAS_ONE_BASE)?1:0;

    /* ----------- 工作区：先把 B 拷到 X ---------------- */
    for(int j=0;j<nrhs;++j)
        std::copy(B + j*ldb, B + j*ldb + m,    X + j*ldx);

    /* ================= 前向  (LOWER) ================= */
    if (uplo == OPENSPBLAS_LOWER)
    {
        for(int colA = 0; colA < m; ++colA)
        {
            /* 找到对角，并先扣除行内已知列贡献 */
            T diag{};
            for(int p=csc_col_ptr[colA]-base; p<csc_col_ptr[colA+1]-base; ++p){
                int row = csc_row_ind[p]-base;
                if(row==colA) { diag = csc_vals[p]; continue; }
                if(row< colA){         // 仅严格下
                    for(int j=0;j<nrhs;++j)
                        X[colA + j*ldx] -= csc_vals[p] * X[row + j*ldx];
                }
            }
            if(abs_val(diag)==T{}) return OPENSPBLAS_STATUS_SINGULAR;
            for(int j=0;j<nrhs;++j)        // 归一化
                X[colA + j*ldx] /= diag;

            /* 用新得到的 x[colA,*] 更新下面各行 */
            for(int p=csc_col_ptr[colA]-base; p<csc_col_ptr[colA+1]-base; ++p){
                int row = csc_row_ind[p]-base;
                if(row > colA){
                    for(int j=0;j<nrhs;++j)
                        X[row + j*ldx] -= csc_vals[p] * X[colA + j*ldx];
                }
            }
        }
        return OPENSPBLAS_STATUS_SUCCESS;
    }

    /* ================= 后向  (UPPER) ================= */
    if (uplo == OPENSPBLAS_UPPER)
    {
        for(int colA = m-1; colA >= 0; --colA)
        {
            T diag{};
            for(int p=csc_col_ptr[colA]-base; p<csc_col_ptr[colA+1]-base; ++p){
                int row = csc_row_ind[p]-base;
                if(row==colA){ diag = csc_vals[p]; continue; }
                if(row> colA){                         // 上三角严格部分
                    for(int j=0;j<nrhs;++j)
                        X[colA + j*ldx] -= csc_vals[p] * X[row + j*ldx];
                }
            }
            if(abs_val(diag)==T{}) return OPENSPBLAS_STATUS_SINGULAR;
            for(int j=0;j<nrhs;++j)
                X[colA + j*ldx] /= diag;

            for(int p=csc_col_ptr[colA]-base; p<csc_col_ptr[colA+1]-base; ++p){
                int row = csc_row_ind[p]-base;
                if(row < colA){
                    for(int j=0;j<nrhs;++j)
                        X[row + j*ldx] -= csc_vals[p] * X[colA + j*ldx];
                }
            }
        }
        return OPENSPBLAS_STATUS_SUCCESS;
    }

    return OPENSPBLAS_STATUS_INVALID;
}



OpenSpB_status OpenSpB_ssptrsm(int m,int nrhs,
    const int*cp,const int*ri,const float*v,
    float*X,int ldx,const float*B,int ldb,
    OpenSpB_base_type base,OpenSpB_uplo_type uplo)
{
    return OpenSpB_spsm<float>(m,nrhs,cp,ri,v,X,ldx,B,ldb,base,uplo);
}
OpenSpB_status OpenSpB_dsptrsm(int m,int nrhs,
    const int*cp,const int*ri,const double*v,
    double*X,int ldx,const double*B,int ldb,
    OpenSpB_base_type base,OpenSpB_uplo_type uplo)
{
    return OpenSpB_spsm<double>(m,nrhs,cp,ri,v,X,ldx,B,ldb,base,uplo);
}
OpenSpB_status OpenSpB_csptrsm(int m,int nrhs,
    const int*cp,const int*ri,const void*v,
    void*X,int ldx,const void*B,int ldb,
    OpenSpB_base_type base,OpenSpB_uplo_type uplo)
{
    return OpenSpB_spsm<std::complex<float>>(m,nrhs,
        cp,ri,
        static_cast<const std::complex<float>*>(v),
        static_cast<std::complex<float>*>(X),ldx,
        static_cast<const std::complex<float>*>(B),ldb,
        base,uplo);
}
OpenSpB_status OpenSpB_zsptrsm(int m,int nrhs,
    const int*cp,const int*ri,const void*v,
    void*X,int ldx,const void*B,int ldb,
    OpenSpB_base_type base,OpenSpB_uplo_type uplo)
{
    return OpenSpB_spsm<std::complex<double>>(m,nrhs,
        cp,ri,
        static_cast<const std::complex<double>*>(v),
        static_cast<std::complex<double>*>(X),ldx,
        static_cast<const std::complex<double>*>(B),ldb,
        base,uplo);
}
